#%%
### This script was written by Eric Deal, 06.22.2022
### Script should be run with python 3
import numpy as np, pandas as pd
grains = ['dn', 'jo', 'mb', 'mk', 'pg', 'sf']
grains_new = {'dn': 'tg', 'jo': 'ng2', 'mb': 'ng3', 'mk': 'ng4', 'pg': 'ng5', 'sf': 'sf'}

# function to calculate settling velocity of sphere as per Dietrich et al,
def R1(logdx):
    # a, b, c, d, f = -3.76715, 1.92944, -0.09815, -0.00575, 0.00056
    a, b, c, d, f = -3.815642429922084, 1.9459392605305115, -0.09016149260089677, -0.008555578212166149, 0.000757318175522988
    return (a + b*logdx + c*logdx**2 + d*logdx**3 + f*logdx**4)

rof = 998.21
nu = 1e-6
g = 9.81

df_density = pd.read_csv('1_var_density_2022.csv').set_index('Unnamed: 0').rename_axis(None)
df_shape = pd.read_csv('2_var_shape_2022.csv').set_index('Unnamed: 0').rename_axis(None)

grains = df_density.index.values.astype('unicode')
dx_veqs = {grain: (df_density['total_grain_ro'][grain] - rof) * g * df_shape.dn[grain]**3 / (rof*nu*nu) for grain in grains}
ws_sphere = {grain: (10**R1(np.log10(dx_veqs[grain])) * (df_density['total_grain_ro'][grain] - rof)*g*nu/rof)**(1./3.) for grain in grains}

dx_veqs = {grain: (df_density['total_grain_ro'][grain] + df_density['dtotal_grain_ro'][grain] - rof) * g * df_shape.dn[grain]**3 / (rof*nu*nu) for grain in grains}
dws_sphere1 = {grain: (10**R1(np.log10(dx_veqs[grain])) * (df_density['total_grain_ro'][grain] + df_density['dtotal_grain_ro'][grain] - rof)*g*nu/rof)**(1./3.) for grain in grains}

dx_veqs = {grain: (df_density['total_grain_ro'][grain] - df_density['dtotal_grain_ro'][grain] - rof) * g * df_shape.dn[grain]**3 / (rof*nu*nu) for grain in grains}
dws_sphere2 = {grain: (10**R1(np.log10(dx_veqs[grain])) * (df_density['total_grain_ro'][grain] - df_density['dtotal_grain_ro'][grain] - rof)*g*nu/rof)**(1./3.) for grain in grains}

# dws_sphere = {grain: dws_sphere1[grain] - dws_sphere2[grain] for grain in grains}
dws_sphere = {grain: 0 for grain in grains}

df = pd.DataFrame({'ws_sphere': ws_sphere, 'dws_sphere': dws_sphere})

df['CD_VES'] = np.nan
df['Re'] = np.nan
for grain in df.ws_sphere.index:
        df['CD_VES'][grain] = 4*(df_density.total_grain_ro[grain]/rof - 1)*g*df_shape.dn[grain] / (3*df.ws_sphere[grain]**2)
        df['Re'][grain] = 1e6*df_shape.A[grain]

df.to_csv('5_var_VES_settling_velocity_2022.csv')
# %%
